#/usr/bin/python
# -*- coding: utf-8 -*-
'''

This module contains classes for load and boundary condition management. 

Created on 7 mars 2012

@author: Ludovic Barrière and Steven Marguet, Institut Clément Ader, 135 av. de Rangueil, 31077 Toulouse, France

© copyright:
    
    TODO
'''
import sys
sys.path.append('../../../icafe')
import parameters

from collections import namedtuple
from domain.domain import GlobalDomain
from analysis.analysis import LinearStatic 
from analysis.strategy import FullModel
import scipy.sparse as sp
from boundaryConditions import BC
from loadConditions import LC

FLOAT = "float64"

#Step = namedtuple("Step", ["NAME", "BC", "BCSTATE", "LC", "LCSTATE", "TYPE", "INCREMENT_PARAMETERS"] )
IncrementParameters = namedtuple("IncParam", ["incrementalMethod", "incrementSize"])


class StepManager(object):
    '''
    
    '''
    
    def __init__(self,analysis=None):
        '''
        '''
        # steps:
        self.steps = [] # [{"NAME": ,"BC": , "BCSTATE": ,"LC": , "LCSTATE": , "TYPE": , "INCREMENT_PARAMETERS": }]
        self.stepsMap = {"U2C": {}, "C2U": []}
        
        self.currentStep = -1
        
        # conditions
        self.bcMap = {"U2C": {}, "C2U": []}
        self.bc = []
        self.lcMap = {"U2C": {}, "C2U": []}
        self.lc = []
        
        # loadcases:
        self.loadCases = {}
        
        if not analysis:
            analysis = LinearStatic
            
        self.default_analysis = analysis
#        self.currentBC = {"ACTIVATED": [], "PROPAGATED": [], "DESACTIVATED": []}
#        self.currentLC = {"ACTIVATED": [], "PROPAGATED": [], "DESACTIVATED": []}
        
        
#        self.disp = {"activated": None, "propagated": None}
#        self.disp = 
#        self.forces = {"activated": None, "propagated": None}
      
    nextstep_kwargs = {"domain": None, "nodeSetMap": None}
    
    def __repr__(self):
        
        m = "Instance of StepManager()\n"
        m+= ">> STEPS:\n"
        for step in self.steps:
            m += step.__repr__() + '\n'
        m+= "\n>> Boundary conditions:\n"
        for bc in self.bc:
            m+= repr(bc) + '\n'
        m+= "\n>> Load conditions:\n"
        for lc in self.lc:
            m+=repr(lc) + '\n'
        m+= "\n>> Loadcases:\n"
        for ldc in self.loadCases.values():
            m+= repr(ldc) + '\n'
            
        return m
        
    def __str__(self):
        
        return self.__repr__()
             
    def addStep(self,stepname,analysis=None, incrementManager=None, solvingStrategy=None):
        '''
        '''
        if not analysis:
            analysis = self.default_analysis
        
        print 'add step:'    
        
        step = Step(stepname,analysis,incrementManager,solvingStrategy)
        
        self.steps.append(step)
        self.stepsMap["U2C"][step.UId]=len(self.steps)-1
        self.stepsMap["C2U"].append(step.UId)
    
    def rmStep(self,stepUId):
        
        stpindex = self.stepsMap['U2C'][stepUId]
        
        self.steps.pop(stpindex)
        del self.stepsMap['U2C'][stepUId]
        self.stepsMap['C2U'].pop(stpindex)
    
    def addLoadCase(self,loadcase,step=-1):
        '''
        '''
        if step:
            self.steps[self.stepsMap["U2C"][step]].addLoadCase(loadcase)
        
        self.loadCases[loadcase.UId]=loadcase
    
    def setNextStep(self, domain):
        '''
            Sets the next step. 
            @param domain: instance of class Domain. Domain concerned by the step.
        '''
#        self.currentBC = {"ACTIVATED": [], "PROPAGATED": [], "DESACTIVATED": []}
#        
#        self.currentLC = {"ACTIVATED": [], "PROPAGATED": [], "DESACTIVATED": []}
        self.currentStep += 1 
        if self.currentStep>=len(self.steps):
            output = None
        else:
            output = self.setStep(self.currentStep, domain)
            

        
        return output
        
        
            
    def setStep(self, CId, domain):      
        '''
            Sets the step identified by a code identifier in a given domain.
            @param CId: int. Code-identifier of the step to set.
            @param domain: instance of class Domain. Domain concerned by the step.
        '''
        
        # Initialize prescribed displacement
#        for i, bcCId in enumerate(self.steps[stepCId]["BC"]):
#            
#            state = self.steps[stepCId]["BCSTATE"][i]
#            
#            # INACTIVE BC:
#            if state == "INACTIVE":
#                continue
#    
#            try:
#                self.currentBC[state].append(bcCId)
#            except KeyError:
#                return KeyError("BC or LC cannot be " + str(state))
#        
#        # Initialize imposed forces
#        for i, lcCId in enumerate(self.steps[stepCId]["LC"]):
#            
#            state = self.steps[stepCId]["LCSTATE"][i]
#            
#            # INACTIVE LC:
#            if state == "INACTIVE":
#                continue
#    
#            try:
#                self.currentLC[state].append(lcCId)
#            except KeyError:
#                return KeyError("BC or LC cannot be " + str(state))
#        
#        try: 
#            dom = nextstep_kwargs["domain"]
#            if isinstance(dom, GlobalDomain()):
#                out = []
#                for subdom in dom.domains:
#                    out.append(self.setStep_domain(stepCId, subdom))
#                return out
#            else:
#                return self.setStep_domain(stepCId, dom)
#        except:
#            pass
#            
#        try:
#            nodeSetMap = nextstep_kwargs["nodeSetMap"]
#            return self.setStep_mesh(stepCId, nodeSetMap)
#        except:
#            pass
        # Associate BC/LC with domains and meshes.
        
        step = self.steps[CId]
        
        domain.initStep(step.UId)
        
        for loadcase in step.loadcases:
        # Adds BC and LC in domain.
            for i,bcId in enumerate(loadcase.bcIds):
                bc = self.bc[self.bcMap["U2C"][bcId]]
                bc.state = loadcase.bcStates[i]
                domain.addBC(bc)
                
        # Adds load cases in domain.
            for i, loadId in enumerate(loadcase.loadIds):
                load = self.lc[self.lcMap["U2C"][loadId]]
                load.state = loadcase.loadStates[i]
                domain.addLoad(load)
        
            domain.addLoadCase(loadcase)
        
        domain.setConstraints()
        
        domain.setAnalysis(step.analysis)
        
        domain.setIncrementMngr(step.incrementMngr)
        
        domain.setSolvingStrategy(step.solvingStrategy)
        
        return step


    def setStepProperties(self, stepUId, analysis=None, incrementManager=None, solvingStrategy=None):
        '''
        
        '''
        step = self.steps[self.stepsMap['U2C'][stepUId]]
        
        print 'set step properties...'
        
        if analysis:
            step.setAnalysis(_analysis)
        if incrementManager:
            step.setIncrementMngr(incrementManager)
        if solvingStrategy:
            step.setSolvingStrategy(solvingStrategy)
            
        return 0
    
    def setStepAnalysis(self, step, _analysis):
        '''
        
        '''
        self.steps[self.stepsMap['U2C'][step]].setAnalysis(_analysis) 
        
        return 0
    
    CONDITION_STATE = ("ACTIVATED", "PROPAGATED", "INACTIVE")
    
    def setConditionState(self, step, conditionCId, conditionType,  STATE = "ACTIVATED" ):
        '''
        
        '''
        
        if isinstance(conditionCId, BoundaryCondition):
            TYPE = "BC"
            
        elif isinstance(conditionCId, LoadCondition):
            TYPE = "LC"
        else:
            return TypeError("Condition must be an instance of classes boundaryCondition() or loadCondition()")
        
        if not STATE in StepManager.CONDITION_STATE:
            return ValueError("Condition state must be in " + repr(StepManager.CONDITION_STATE))
        
        
        
        return 0

    def addCondition(self, condition):
        
        if isinstance(condition, BC):
            TYPE = "BC"
            self.bc.append(condition)
            nbc = len(self.bc)
            self.bcMap["U2C"][condition.UId]=nbc-1
            self.bcMap["C2U"].append(condition.UId)
            
        elif isinstance(condition, LC):
            TYPE = "LC"
            self.lc.append(condition)
            nlc = len(self.lc)
            self.lcMap["U2C"][condition.UId]=nlc-1
            self.lcMap["C2U"].append(condition.UId)
            
        else:
            return TypeError("Condition must be an instance of classes boundaryCondition() or loadCondition()")
        
        
        return 0
    
    def rmCondition(self, condition):
        '''
        
        '''
        if isinstance(condition, BoundaryCondition):
            TYPE = "BC"
        elif isinstance(condition, LoadCondition):
            TYPE = "LC"
        else:
            return TypeError("Condition must be an instance of classes boundaryCondition() or loadCondition()")
        
        
        return 0
    
    def startStepAnalysis(self,domain):
        '''
            
        '''
        step = self.steps[self.currentStep]
        
        # initialize analysis:
        analysis = step.analysis(step.incrementManager, step.solvingStrategy)
        
        # start analysis of domain:
        return analysis.start(domain)
    
    
    def manageSteps(self):
        '''
            Calls the user interface which enables to organize steps and set parameters and load cases.
        '''
        pass
        
    
class Step():
    '''
        Contains informations about a step: load cases, type of analysis, type of solving strategy and type of 
        incremental procedure.
    '''
    def __init__(self,UId,_analysis, incrementManager = None, solvingStrategy = None):
        '''
            Constructor of class Step:
            @param UId: str or int, user-identifier of the step.
            @param analysis: Class of the analysis module, corresponding to the chosen analysis
            @param incrementManager=IncrementManager(1): instance of IncrementManager class, optional. default value is 
            a standard increment manager with increment size set to 1.
            @param solvingStrategy: Class, optional. Default class is the FullModel solving strategy.
        '''
        self.UId = UId
        self.loadcases = []
        self.analysis = _analysis()
        if incrementManager:
            self.incrementMngr = incrementManager
        else:
            self.incrementMngr = IncrementManager(nincr=1, stepId = self.UId)
        if solvingStrategy:
            self.solvingStrategy = solvingStrategy
        else:    
            self.solvingStrategy = FullModel()
    
    def __repr__(self):
        
        m = "step " + str(self.UId) + ":\n"
        m+= " - loadcases: \n"
        for lc in self.loadcases:
            m +="  * " + str(lc.UId) + '\n'
        m+= " - analysis: " + repr(self.analysis) + '\n'
        m+= " - incrementManager: " + self.incrementMngr.__repr__()+ '\n'
        m+= " - solvingStrategy: " + self.solvingStrategy.__repr__()+ '\n'
        
        return m
    
    def addLoadCase(self,loadcase):
        '''
            Adds a load case to the definition of the step.
            @param loadcase: instance of the LoadCase class. 
        '''
        self.loadcases.append(loadcase)
    
    def rmLoadCase(self, loadcase):
        '''
            Removes a load case from the definition of the step.
            @param loadcase: instance of the LoadCase class.
        '''
        self.loadcases.remove(loadcase)
        
    
    def setAnalysis(self,analysis):
        '''
            @param analysis: Class of the analysis module.
        '''
        self.analysis = analysis()   
        
    def setIncrementMngr(self, incrementManager):
        '''
            @param incrementManager: Instance of IncrementManager class.
        '''
        self.incrementMngr = incrementManager

    def setSolvingStrategy(self, solvingStrategy):
        '''
            @param solvingStrategy: Class of the strategy module.
        '''
        self.solvingStrategy = solvingStrategy


#-------------------------------------------------------------------------------------------------------#
#                                             INCREMENT MANAGERS
#-------------------------------------------------------------------------------------------------------#

class IncrementManager():
    '''
        Default class of Increment Manager. Manages the incremental procedure.
        Basically displacement and load factors are defined in [0,1] and incremented by the "setNextIncr" method.
        
    '''
    def __init__(self, disp=0, load=0, nincr=1, maxdisp=1, maxload=1, sequence=None, stepId='0'):
        '''
        @param disp=0: initial value of the displacement factor.
        @param load=0: initial value of the load factor.
        @param delta=1: default value of the size of the increment.
        '''
        self.stepId = stepId
        self.incr = 0
        self.nincr = nincr
        
        
        self.disp  = 0  # total imposed displacement factor  
        self.load  = 0  # total imposed load factor
        self.dload = 0  # incremental factor of load
        
        self.maxdisp = maxdisp # maximum level of prescribed displacement
        self.maxload = maxload # maximum level of prescribed loads
        self.iter  = 0
        self.sequence = sequence
        
        self.updateLhs = True
        
        if sequence:
            self.nincr = len(sequence)
            try:
                self.delta = self.sequence.__iter__()
            except:
                raise TypeError('Incremental sequence must be a list of increment size')
        else:
            self.sequence = [self.maxdisp/float(nincr)]*nincr
            self.delta = self.sequence.__iter__()
        
    def __repr__(self):
        
        return 'basic increment manager with ' + str(self.nincr) + ' initial increments.' 
    
    def getStateKey(self, level='iter'):
    
        keydict = {'step':0,
                   'incr':1,
                   'iter':2,
                   }
        keylist = ['ST','-IN','-IT']
        levels = [str(self.stepId),str(self.incr),str(self.iter)]
        key = ''
        for i in range(keydict[level]+1):
            key += keylist[i] + levels[i]
            
        return key
        
    def setNextIncr(self, **nextinc_kwargs):
        '''
            Increments the displacement and load factors.
        '''
        if self.load>=self.maxload and self.disp>=self.maxdisp:
            print '-> max load reached'
            return False, False
        else:
            # Calculate displacement and load increments:
            try:
                d = self.delta.next()
            except StopIteration:
                d = self.sequence[-1]
                
            self.disp = min(self.disp+d, self.maxdisp)
            self.load = min(self.load+d, self.maxload)
            
            self.iter = 1
        
        self.incr += 1
         
        return self.disp, self.load
    
    def setNextIter(self, *args):
        '''
            Set next iteration and update the displacement and load factors during an iterative procedure.
        '''
        self.iter += 1
        
        # can modify ddisp and dload during increment
        
        return self.load
        
    def cutNextIncrement(self, cutfactor=0.5):
        '''
            Cut the increment. Multiplies by the "cutfactor".
            @param cutfactor=0.5
        '''
        self.delta *= cutfactor
        
        return 0
    
    def cutRestartIncrement(self, cutfactor=0.5):
        '''
            Cut and restart the increment.
            @param cutfactor=0.5
        '''
        self.disp -= self.delta
        self.load -= self.delta
        
        self.cutNextIncrement(cutfactor)
        
        return 0 

def IncrementDecorator(IncrementManager):
    '''
        Decorator of the IncrementManager class.
    '''
    def __init__(self, incrementManager):
        disp = incrementManager.disp
        load = incrementManager.load
        delta = incrementManager.delta
        super(incrementDecorator,self).__init__(disp,load,delta)

def ArcLengthIncrement(IncrementDecorator):
    '''
        Increment manager which relies on the Arc-Length path-following technique [M.-A. Crisfield, 1991].
    '''
    def __init__(self, incrementManager):
        
        super(ArcLengthIncrement,self).__init__(incrementManager)
    
    def __repr__(self):
        
        return 'arc-length increment manager'
    # redefine methods:
    def setNextIncr():
        
        pass
    
    def setNextIter():
        
        pass
    
def RiksIncrement(IncrementDecorator):
    '''
        Increment manager which relies on the Riks path-following technique [M.-A. Crisfield, 1991].
    '''
    def __init__(self, incrementManager):
        
        super(RiksIncrement,self).__init__(incrementManager)

    def __repr__(self):
        
        return 'Riks increment manager'    
    # redefine methods:
    def setNextIncr():
        
        pass
    
    def setNextIter():
        
        pass
#-------------------------------------------------------------------------------------------------------#
#                                             LOADCASES
#-------------------------------------------------------------------------------------------------------#

class LoadCase():
    '''
        Contains load case definitions (association of boundary and load conditions).
    '''
    def __init__(self, UId, bcIds=[], loadIds=[], bcStates=None, loadStates=None):
        '''
            @param UId: User identifier of the load case.
            @param bcIds=[]: list of identifiers of the boundary conditions of the load case.
            @param lcIds=[]: list of identifiers of the load conditions of the load case.
            @param loadState
        '''
        self.UId = UId
        self.bcIds = bcIds
        self.loadIds = loadIds
        self.bcStates=[]
        self.loadStates=[]
        if not bcStates:
            self.bcStates = ["ACTIVATED"]*len(bcIds)
        else:
            self.bcStates = bcStates
        if not loadStates:
            self.loadStates = ["ACTIVATED"]*len(loadIds)
        else:
            self.loadStates = loadStates 
    
    def __repr__(self):
        m= "loadcase "+ str(self.UId)+ ":\n"
        m+= " - BC: "+ str(self.bcIds) + "\n"
        m+= " - LC: "+ str(self.loadIds)
        
        return m
       
    def addCondition(self,condition, state="ACTIVATED"):
        '''
            Adds condition to the load case.
            @param condition: Instance of the boundary or load condition class
            @param state="ACTIVATED": str, optional. State of the condition in the load case ("ACTIVATED", "PROPAGATED" or "INACTIVE")
        '''
        if isinstance(condition,BC):
            self.bcIds.append(condition.UId)
            self.bcStates.append(state)
        elif isinstance(condition,LC):
            self.loadIds.append(condition.UId)  
            self.loadStates.append(state)
        
    def rmCondition(self,condition):
        '''
            Adds condition to the load case.
            @param condition: Instance of the boundary or load condition class
        '''
        if isinstance(condition,BC):
            self.bcIds.remove(condition.UId)
        elif isinstance(condition,LC):
            self.loadIds.remove(condition.UId)  

#class Step(object):
#    
#    def __init__(self):
#        self.bc = []
#        self.lc = []         
#    
    
# BC
#class BoundaryCondition(object):     # BC Parent Class
#
#    def __init__(self, nodeSetUIds = []):
#        self.nodeSetUIds = nodeSetUIds
#    
#    def __repr__(self):
#        val = "** Boundary condition properties:\n"
#        val += "* Node sets: "+repr(self.nodeSetUIds)+"\n"
#        return val
#
#
#class StandardBc(BoundaryCondition):    # BC Child Class
#
#    def __init__(self, nodeSetUIds=[], constrainedDofs=[], imposedValues=[]):
#        BoundaryCondition.__init__(self, nodeSetUIds)
#        self.constrainedDofs =  constrainedDofs
#        self.values = imposedValues # [ values of dofs ] or values_of_dofs(nodeCoo)
#        # ...
#    
#    def __repr__(self):
#        val = BoundaryCondition.__repr__(self)
#        val += "* Constrained degrees of freedom: "+repr(self.constrainedDofs)+"\n"
#        val += "* Imposed values: "+repr(self.values)+"\n"
#        return val
#    # ...
#
#    
## LOADS    
#class LoadCondition(object):        # LC Parent Class
#
#    def __init__(self, ):
#        self.nodeSets = []
#    
#    def __repr__(self):
#        val = "** Load condition properties:\n"
#        val += "* Node sets: "+repr(self.nodeSets)+"\n"
#        return val
#    
#class DistributedLoad(LoadCondition):    # LC Child Class
#
#    def __init__(self, ):
#        LoadCondition.__init__(self)
#        # ...
#
#class StandardLoad(LoadCondition):    # LC Child Class
#
#    def __init__(self, nodeSetUIds=[], constrainedDuals=[], imposedValues=[]):
#        LoadCondition.__init__(self, nodeSetUIds)
#        self.constrainedDuals = list(constrainedDuals)
#        self.values = list(imposedValues)
#        # ...
#    
#    def __repr__(self):
#        val = LoadCondition.__repr__(self)
#        val += "* Constrained duals: "+repr(self.constrainedDuals)+"\n"
#        val += "* Associated values: "+repr(self.values)+"\n"
#        return val
#    